import argparse
import itertools
import os
import pprint
import sqlite3
import sys
import json
import collections
import pysnmp.hlapi as hlapi
import pysnmp.proto.rfc1902 as rfc1902
from pysnmp.hlapi import *
import re
import socket  # 添加 socket 模块的导入

# OIDs  - 这些应该通常有效， 但可能需要调整
IF_DESCR_OID = '.1.3.6.1.2.1.2.2.1.2'
IF_PHYS_ADDRESS_OID = '.1.3.6.1.2.1.2.2.1.6'
SYS_NAME_OID = '.1.3.6.1.2.1.1.5.0'

def snmp_get(host, oid, community='public@123'):
    """使用 SNMP GET 获取单个 OID 的值."""
    try:
        for (errorIndication,
             errorStatus,
             errorIndex,
             varBinds) in hlapi.getCmd(hlapi.SnmpEngine(),
                                        hlapi.CommunityData(community),
                                        hlapi.UdpTransportTarget(('192.168.2.1', 161), timeout=4.0, retries=3), #尝试使用 192.168.2.1，而非host
                                        hlapi.ContextData(),
                                        hlapi.ObjectType(hlapi.ObjectIdentity(oid)),
                                        lookupMib=False):

            if errorIndication:
                print(f'SNMP GET error: {errorIndication}')
                return None
            elif errorStatus:
                print(f'SNMP GET error: {errorStatus} at {errorIndex}')
                return None
            else:
                for varBind in varBinds:
                    return str(varBind[1])  # 返回字符串值
        return None
    except socket.gaierror as e:
        print(f"网络错误: 无法解析主机名 192.168.2.1: {e}")#主机名改为IP地址
        return None

def snmp_walk(host, oid, format='str', strip_prefix=True, community='public@123'):  # modify
    res = []
    for (errorIndication,
         errorStatus,
         errorIndex,
         varBinds) in hlapi.nextCmd(hlapi.SnmpEngine(),
                                    hlapi.CommunityData(community),
                                    hlapi.UdpTransportTarget(('192.168.2.1', 161), timeout=4.0, retries=3), #尝试使用 192.168.2.1，而非host
                                    hlapi.ContextData(),
                                    hlapi.ObjectType(hlapi.ObjectIdentity(oid)),
                                    lookupMib=False,
                                    lexicographicMode=False):
        if errorIndication:
            print(f'SNMP error: "{str(errorIndication)}". Status={str(errorStatus)}')
            return {}
        elif errorStatus:
            print('errorStatus: %s at %s' % (errorStatus.prettyPrint(),
                                                             errorIndex and varBinds[int(errorIndex) - 1][0] or '?'))
            return {}
        else:
            for x in varBinds:
                k, v = x
                if strip_prefix:
                    k = str(k)[len(str(oid)) + 1:]
                if isinstance(v, rfc1902.Integer):
                    res.append((str(k), int(v)))
                else:
                    if format == 'numbers':
                        res.append((str(k), v.asNumbers()))
                    elif format == 'hex':
                        res.append((str(k), v.asOctets().hex()))
                    elif format == 'raw':
                        res.append((str(k), v))
                    elif format == 'bin':
                        res.append((str(k), v.asOctets()))
                    elif format == 'int':
                        res.append((str(k), int(v)))
                    elif format == 'preview':
                        res.append((str(k), str(v)))
                    elif format == 'any':
                        try:
                            res.append((str(k), v.asOctets().decode('utf-8')))
                        except UnicodeDecodeError:
                            res.append((str(k), '0x' + v.asOctets().hex()))
                    elif format == 'str':
                        res.append((str(k), v.asOctets().decode(v.encoding)))
                    else:
                        assert False, "Unknown format for walk()."
    res = {a: b for a, b in res}
    return res

def main():

    parser = argparse.ArgumentParser(description="使用SNMP获取网络拓扑信息")
    parser.add_argument("-f", "--file", type=str, required=False, help="包含交换机信息的文件路径")
    args = parser.parse_args()

    if args.file: #如果传入文件，从文件获取交换机信息
      try:
        with open(args.file, 'r') as f:
            line = f.readline().strip()
            community, snmp_port, ip = line.split(',')
      except Exception as e:
        print(f"解析文件失败,请检查文件格式是否正确! 错误信息:{e}")
        return # 终止程序
    else: #从命令行获取交换机信息
        community = "public@123"
        snmp_port = 161
        ip = "192.168.2.1"
    # 获取交换机名称
    switch_name = snmp_get(ip, SYS_NAME_OID, community)
    if not switch_name:
        print("无法获取交换机名称, 请检查SNMP配置或OID是否正确")
        switch_name = "Unknown" #可以自定义名称
    print(f"交换机名称: {switch_name}")

    # 接口描述
    if_descr = snmp_walk(ip, IF_DESCR_OID, format='str', community=community)
    # 物理地址
    if_phys_address = snmp_walk(ip, IF_PHYS_ADDRESS_OID, format='hex', community=community) #format 修改为 hex

    # Prepare data for vis.js
    nodes = []
    edges = []

    # 添加交换机节点
    nodes.append({'id': switch_name, 'label': switch_name})

    # 添加接口节点和边
    if if_descr and if_phys_address:  # 确保数据不为空
        for if_index, if_desc in if_descr.items():
            phys_address = if_phys_address.get(if_index, "Unknown") #从物理地址信息中，获取对应物理地址，如果没有就显示Unknown
            node_id = f"{switch_name}-{if_index}" #接口添加ID
            nodes.append({'id':node_id, 'label': f"{if_desc}\nMAC: {phys_address}"})
            edges.append({'from': switch_name, 'to': node_id, 'label': "连接"})  #所有节点默认连接到交换机

    # Convert to vis.js format
    nodes_list = [dict(node) for node in nodes]
    edges_list = [dict(edge) for edge in edges]

    # 写入 nodes.js
    with open("nodes.js", "w") as f:
        f.write("var nodes = new vis.DataSet(" + json.dumps(nodes_list, indent=4) + ");")
    # 写入 edges.js
    with open("edges.js", "w") as f:
        f.write("var edges = new vis.DataSet(" + json.dumps(edges_list, indent=4) + ");")
    print("生成拓扑文件成功！请打开 snmplldpnetworkmap.html 查看")

if __name__ == '__main__':
  main()